#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : mngs_download_genomes.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/11/3 14:35
@Description :
下载 Staphylococcus aureus 完整 RefSeq 基因组
"""
from argparse import ArgumentParser, Namespace
from os import cpu_count
from pathlib import Path
import subprocess
import logging
from multiprocessing import Pool
from typing import Sequence
from datetime import datetime


# --------
# 日志
# --------
logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] %(levelname)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
LOGGER = logging.getLogger("mNGS")


# --------
# 软件
# --------
ARIA2C = Path("/home/bioinfo/software/miniconda3/envs/aria2c/bin/aria2c")
DATASETS = Path("/data/jinlj/software/jenv/bin/datasets")
PYTHON3 = Path("/home/bioinfo/software/miniconda3/envs/python-venv/bin/python3")
JIN_PYTHON3 = Path("/data/jinlj/software/jenv/bin/python3")
BWA = Path("/home/bioinfo/software/miniconda3/envs/bwa/bin")


# --------
# 脚本
# --------
SCRIPTS = Path(__file__).resolve().parent
PARSE_JSONL_PY =  SCRIPTS / "mngs_parse_jsonl.py"
REPLACE_REF_PY = SCRIPTS / "mngs_replace_ref.py"
KMER_PROCESS_PY = Path("/data/jinlj/project/panGenome/bin/step1_kmer_process.py")
MERGE_PY = Path("/data/jinlj/project/panGenome/bin/step2_merge.py")


# --------
# 辅助函数
# --------
def _run_or_die(sh_cmd: str, sh_file: Path, log_file: Path, ok_file: Path):
    """
    执行 shell 命令。若失败则记录错误，成功则创建 .ok 文件。
    """
    if ok_file.exists():
        LOGGER.info(f"跳过 {sh_file}（已完成）")
        return

    try:
        LOGGER.info(f"执行任务：{sh_file}")
        with open(sh_file, "w", encoding="utf-8") as fs:
            fs.write(sh_cmd)

        with open(log_file, "w", encoding="utf-8") as fl:
            result = subprocess.run(
                f"bash {sh_file}",
                shell=True,
                text=True,
                stdout=fl,
                stderr=fl,
            )

        if result.returncode != 0:
            LOGGER.error(f"执行失败：{sh_file}，退出码 {result.returncode}")
            return

        ok_file.touch()
        LOGGER.info(f"完成：{sh_file}")

    except Exception as e:
        LOGGER.error(f"执行出错：{sh_file}，错误：{e}")


def _run_cmds(sh_files: list[Path], sh_cmds: list[str], cpu_n: int = 1, check: bool = True):
    """
    并行执行多个命令脚本。
    """
    ok_files = [f.with_suffix(".ok") for f in sh_files]
    log_files = [f.with_suffix(".log") for f in sh_files]

    with Pool(cpu_n) as pool:
        for sh_cmd, sh_file, ok_file, log_file in zip(sh_cmds, sh_files, ok_files, log_files):
            pool.apply_async(_run_or_die, args=(sh_cmd, sh_file, log_file, ok_file, ))
        pool.close()
        pool.join()

    if check:
        for ok_file in ok_files:
            if not ok_file.exists():
                raise RuntimeError(f"缺少完成标志文件：{ok_file}")


def _check_out_files(out_files: Sequence[Path]) -> None:
    """
    检查输出文件是否存在或为空，如果为空直接抛异常。

    :param out_files: 输出文件列表
    :raises RuntimeError: 如果文件不存在或为空
    """
    missing_files = [f for f in out_files if not f.exists()]
    empty_files = [f for f in out_files if f.exists() and f.stat().st_size == 0]

    if missing_files or empty_files:
        messages = []
        if missing_files:
            messages.append(f"不存在文件: {', '.join(map(str, missing_files))}")
        if empty_files:
            messages.append(f"为空文件: {', '.join(map(str, empty_files))}")
        raise RuntimeError("; ".join(messages))


# --------
# 分析流程
# --------
class Begin:
    """
    启动分析
    """

    def __init__(self, args: Namespace):
        self.step_n = 0
        self.begin_time = datetime.now()
        self.tax_id: int = args.tax_id
        self.out_dir: Path = Path(args.out_dir).absolute()
        self.log_dir: Path = args.out_dir / "log"
        self.all_genome = args.all_genome
        self.ref_id = args.ref_id
        self.pan_id = args.pan_id
        self.pangenome = Path(args.pangenome).absolute()

    def run(self) -> None:
        LOGGER.info(f"开始下载物种[{self.tax_id}]的所有基因组")
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.log_dir.mkdir(parents=True, exist_ok=True)


class Download(Begin):
    """
    下载基因组元信息
    """
    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.download_out_name = f"step{self.step_n}_download"
        self.download_out_dir = self.out_dir / self.download_out_name
        self.download_log_dir = self.log_dir / self.download_out_name
        self.download_out_file = self.download_out_dir / f"{self.tax_id}.metadata.zip"
        self.download_out_jsonl = (
                self.download_out_dir /
                "ncbi_dataset" /
                "data" /
                f"assembly_data_report.jsonl"
        )
        self.download_sh_files = [self.download_log_dir / f"{self.download_out_name}.sh"]
        # --------
        # 判断是否下载所有的基因组
        # --------
        if self.all_genome:
            level_flag = " "
        else:
            level_flag = " --assembly-level complete "
        self.download_sh_cmds = [
            (
                f"#!/bin/bash\n\n"
                f"set -vexuo pipefail\n\n"
                f"{DATASETS} download genome "
                f"taxon {self.tax_id} "
                f"{level_flag}"
                f"--include genome "
                f"--dehydrated "
                f"--filename {self.download_out_file}\n\n"
                f"cd {self.download_out_dir} && unzip -o {self.download_out_file}\n\n"
            )
        ]
        self.download_cpu_n = 1

    def run(self) -> None:
        super().run()
        self.download_out_dir.mkdir(parents=True, exist_ok=True)
        self.download_log_dir.mkdir(parents=True, exist_ok=True)
        _run_cmds(self.download_sh_files, self.download_sh_cmds, self.download_cpu_n)
        _check_out_files([self.download_out_jsonl])


class Parse(Download):
    """
    解析静态文件地址
    """
    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.parse_out_name = f"step{self.step_n}_parse"
        self.parse_out_dir = self.out_dir / self.parse_out_name
        self.parse_log_dir = self.log_dir / self.parse_out_name
        self.parse_out_file = self.parse_out_dir / f"urls.csv"
        self.parse_sh_files = [self.parse_log_dir / f"{self.parse_out_name}.sh"]
        self.parse_sh_cmds = [
            (
                f"#!/bin/bash\n\n"
                f"set -vexuo pipefail\n\n"
                f"{PYTHON3} {PARSE_JSONL_PY} "
                f"--in_file {self.download_out_jsonl} "
                f"--out_file {self.parse_out_file}\n\n"
            )
        ]
        self.parse_cpu_n = 1

    def run(self) -> None:
        super().run()
        self.parse_out_dir.mkdir(parents=True, exist_ok=True)
        self.parse_log_dir.mkdir(parents=True, exist_ok=True)
        _run_cmds(self.parse_sh_files, self.parse_sh_cmds, self.parse_cpu_n)
        _check_out_files([self.parse_out_file])


class Aria2C(Parse):
    """
    批量下载基因组文件
    """
    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.aria2c_out_name = f"step{self.step_n}_aria2c"
        self.aria2c_out_dir = self.out_dir / self.aria2c_out_name
        self.aria2c_log_dir = self.log_dir / self.aria2c_out_name
        self.aria2c_sh_files = [self.aria2c_log_dir / f"{self.aria2c_out_name}.sh"]
        self.aria2c_sh_cmds = [
            (
                f"#!/bin/bash\n\n"
                f"set -vexuo pipefail\n\n"
                f"{ARIA2C} -c -x16 -s16 -j 16 -i {self.parse_out_file} -d {self.aria2c_out_dir}\n\n"
            )
        ]
        self.aria2c_cpu_n = 1

    def run(self) -> None:
        super().run()
        self.aria2c_out_dir.mkdir(parents=True, exist_ok=True)
        self.aria2c_log_dir.mkdir(parents=True, exist_ok=True)
        _run_cmds(self.aria2c_sh_files, self.aria2c_sh_cmds, self.aria2c_cpu_n, check=False)


class CpGenomes(Aria2C):
    """
    创建物种目录并拷贝基因组序列文件
    """

    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.cpg_out_name = f"step{self.step_n}_cp_genomes"
        self.cpg_out_dir = self.out_dir / "data"
        self.cpg_log_dir = self.log_dir / self.cpg_out_name
        self.cpg_in_fns = []
        self.cpg_species = []
        self.cpg_sh_files = []
        self.cpg_sh_cmds = []
        self.cpg_species_dirs = []
        self.cpg_cpu_n = cpu_count() // 2


    def run(self) -> None:
        super().run()
        self.cpg_out_dir.mkdir(parents=True, exist_ok=True)
        self.cpg_log_dir.mkdir(parents=True, exist_ok=True)
        self.cpg_in_fns = list(self.aria2c_out_dir.glob("*.fna.gz"))
        self.cpg_species = [
            f"{in_fn.name.split('_')[0]}_{in_fn.name.split('_')[1]}"
            for in_fn in self.cpg_in_fns
        ]
        self.cpg_sh_files = [
            self.cpg_log_dir / f"{species}.sh"
            for species in self.cpg_species
        ]
        self.cpg_species_dirs = [
            self.cpg_out_dir / species
            for species in self.cpg_species
        ]
        self.cpg_sh_cmds = [
            (
                f"#!/bin/bash\n\n"
                f"set -vexuo pipefail\n\n"
                f"mkdir {species_dir}\n\n"
                f"cp {in_fn} {species_dir}\n\n"
                f"cd {self.cpg_out_dir}\n\n"
                f"gunzip {species_dir / in_fn.name}\n\n"
            )
            for species_dir, in_fn in zip(self.cpg_species_dirs, self.cpg_in_fns)
        ]
        _run_cmds(self.cpg_sh_files, self.cpg_sh_cmds, self.cpg_cpu_n)


class KmerProcess(CpGenomes):
    """
    将基因组切割为kemer
    """

    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.kmers_out_name = f"step{self.step_n}_kmer"
        self.kmers_out_dir = self.out_dir / self.kmers_out_name
        self.kmers_log_dir = self.log_dir / self.kmers_out_name

    def run(self) -> None:
        super().run()
        self.kmers_out_dir.mkdir(parents=True, exist_ok=True)
        self.kmers_log_dir.mkdir(parents=True, exist_ok=True)
        _run_cmds(
            [self.kmers_log_dir / f"{self.kmers_out_name}.sh"],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -vexuo pipefail\n\n"
                    f"{JIN_PYTHON3} {KMER_PROCESS_PY} "
                    f"--data_dir {self.cpg_out_dir} "
                    f"--ref_accession {self.ref_id} "
                    f"--output_dir {self.kmers_out_dir}\n\n"
                )
            ],
            1
        )


class KmerMerge(KmerProcess):
    """
    将kmer拼接到参考基因组
    """

    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.merge_out_name = f"step{self.step_n}_merge"
        self.merge_out_dir = self.out_dir / self.merge_out_name
        self.merge_log_dir = self.log_dir / self.merge_out_name
        self.merge_out_file = self.merge_out_dir / f"{self.pan_id}.merge.fna"
        self.merge_ref_genome = ""

    def run(self) -> None:
        super().run()
        self.merge_log_dir.mkdir(parents=True, exist_ok=True)
        self.merge_out_dir.mkdir(parents=True, exist_ok=True)
        self.merge_ref_genome = list((self.cpg_out_dir / self.ref_id).glob("*.fna"))[0]
        _run_cmds(
            [self.merge_log_dir / f"{self.merge_out_name}.sh"],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -vexuo pipefail\n\n"
                    f"{JIN_PYTHON3} {MERGE_PY} "
                    f"--ref_genome {self.merge_ref_genome} "
                    f"--kmer_dir {self.kmers_out_dir} "
                    f"--output_file {self.merge_out_file}\n\n"
                )
            ],
            1
        )
        _check_out_files([self.merge_out_file])


class ReplaceRef(KmerMerge):
    """
    替换原始数据库中的参考序列
    """

    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.replace_out_name = f"step{self.step_n}_replace"
        self.replace_out_dir = self.out_dir / self.replace_out_name
        self.replace_log_dir = self.log_dir / self.replace_out_name
        self.replace_out_file = self.replace_out_dir / self.pangenome.name

    def run(self) -> None:
        super().run()
        self.replace_out_dir.mkdir(parents=True, exist_ok=True)
        self.replace_log_dir.mkdir(parents=True, exist_ok=True)
        _run_cmds(
            [self.replace_log_dir / f"{self.replace_out_name}.sh"],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -vexuo pipefail\n\n"
                    f"{PYTHON3} {REPLACE_REF_PY} "
                    f"--pan_genome {self.pangenome} "
                    f"--ref_fasta {self.merge_out_file} "
                    f"--out_file {self.replace_out_file} "
                    f"--species_id {self.pan_id}\n\n"
                    f"export PATH={BWA}:$PATH\n\n"
                    f"bwa index {self.replace_out_file}\n\n"
                )
            ],
            1
        )
        _check_out_files([self.replace_out_file])


class End(ReplaceRef):
    """最终步骤：输出耗时"""

    def __init__(self, args: Namespace):
        super().__init__(args)

    def run(self):
        super().run()
        elapsed = datetime.now() - self.begin_time
        LOGGER.info(f"任务完成，总耗时：{elapsed}")


if __name__ == "__main__":
    parser = ArgumentParser(description="批量下载物种所有完整的基因组")
    parser.add_argument("--tax_id", type=int, required=True, default=None, help="物种ID")
    parser.add_argument("--all_genome", action="store_true", default=False, help="是否下载所有的基因组")
    parser.add_argument("--ref_id", type=str,  required=True, default=None, help="用作参考的基因组accession")
    parser.add_argument("--pan_id", type=str, required=True, default=None, help="pan库id")
    parser.add_argument("--out_dir", type=Path, required=True, default=None, help="输出目录")
    parser.add_argument("--pangenome", type=Path, required=True, default=None, help="原始的pan库序列")
    End(args=parser.parse_args()).run()
